# -*- coding: utf-8 -*-
"""
YOLOv5目标检测模块
"""
import cv2
import torch
import numpy as np
from src.config import YOLO_CONFIG, TRACKING_CONFIG


class YOLODetector:
    def __init__(self, model_path=None):
        """初始化YOLOv5模型"""
        # 从config.py读取配置，这些参数已写死但可通过配置文件修改
        self.model_path = model_path or YOLO_CONFIG['model_path']
        self.conf_threshold = YOLO_CONFIG['conf_threshold']
        self.iou_threshold = YOLO_CONFIG['iou_threshold']
        self.device = YOLO_CONFIG['device']
        self.target_classes = TRACKING_CONFIG['target_classes']

        # 加载YOLOv5模型
        self.model = self._load_model()
        self.class_names = self.model.names  # 获取类别名称

        print(f"YOLOv5模型加载成功! 设备: {self.device}")
        print(f"可检测类别: {len(self.class_names)}个")

    def _load_model(self):
        """加载YOLOv5模型"""
        try:
            model = torch.hub.load('ultralytics/yolov5', 'custom',
                                   path=self.model_path,
                                   device=self.device)
            model.conf = self.conf_threshold  # 设置置信度阈值
            model.iou = self.iou_threshold  # 设置IOU阈值
            return model
        except Exception as e:
            print(f"模型加载失败: {e}")
            raise

    def detect(self, frame):
        """
        对输入帧进行目标检测

        参数:
            frame: 输入图像帧 (BGR格式)

        返回:
            detections: 检测结果列表，每个元素为字典包含检测信息
            frame_with_bbox: 绘制了边界框的图像帧 (BGR格式)
        """
        # YOLOv5推理
        results = self.model(frame)

        # 解析结果
        detections = []
        frame_with_bbox = frame.copy()

        # 获取检测到的目标信息
        for *xyxy, conf, cls in results.xyxy[0]:
            class_id = int(cls)
            # 只保留我们感兴趣的类别
            if class_id in self.target_classes:
                detection = {
                    'bbox': [int(coord) for coord in xyxy],  # [x1, y1, x2, y2]
                    'confidence': float(conf),
                    'class_id': class_id,
                    'class_name': self.class_names[class_id]
                }
                detections.append(detection)

                # 在图像上绘制边界框和标签
                if TRACKING_CONFIG['draw_bbox']:
                    self._draw_bounding_box(frame_with_bbox, detection)

        return detections, frame_with_bbox

    def _draw_bounding_box(self, frame, detection):
        """在图像上绘制边界框和标签"""
        x1, y1, x2, y2 = detection['bbox']
        class_name = detection['class_name']
        confidence = detection['confidence']

        # 绘制矩形框
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 绘制标签背景
        label = f"{class_name} {confidence:.2f}"
        label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
        cv2.rectangle(frame, (x1, y1 - label_size[1] - 10),
                      (x1 + label_size[0], y1), (0, 255, 0), -1)

        # 绘制标签文字
        cv2.putText(frame, label, (x1, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)